import torch
import torch.optim as optim
import torch.nn as nn
from torch.utils.data import DataLoader
from data import ImageSimilarityDataset
from model import SimilarityModel
import os

def train_model():
    # 检查是否有可用的GPU
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    print(f"Using device: {device}")

    # 超参数
    num_epochs = 1000
    batch_size = 64
    learning_rate = 0.001
    image_dir = 'output_frames'
    similarity_file = os.path.join(image_dir, "frame_similarity.npy")

    # 数据集和数据加载器
    dataset = ImageSimilarityDataset(image_dir, similarity_file)
    train_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

    # 获取样本图像的形状
    sample_img, _ = dataset[0]
    input_shape = sample_img.shape

    # 模型、损失函数和优化器
    model = SimilarityModel(input_shape=input_shape).to(device)
    criterion = nn.MSELoss()
    optimizer = optim.Adam(model.parameters(), lr=learning_rate)

    # 训练循环
    for epoch in range(num_epochs):
        print('epoch:',epoch)
        running_loss = 0.0
        for i, (img, similarity) in enumerate(train_loader):
            print('i:', i)
            # 将数据移动到GPU
            img, similarity = img.to(device), similarity.to(device)
            
            optimizer.zero_grad()
            outputs = model(img).squeeze()  # 确保输出是一个一维张量
            similarity = similarity.squeeze()  # 确保目标是一个一维张量
            loss = criterion(outputs, similarity)
            loss.backward()
            
            # 梯度裁剪
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            
            optimizer.step()

            running_loss += loss.item()
            if i % 10 == 9:    # 每10个batch打印一次
                print(f"[{epoch + 1}, {i + 1}] loss: {running_loss / 10:.3f}")
                running_loss = 0.0
                torch.save(model.state_dict(), 'similarity_model.pth')

    # 保存模型
    torch.save(model.state_dict(), 'similarity_model.pth')
    print('Finished Training and Model Saved')

if __name__ == '__main__':
    train_model()
